package com.atguigu0.streaming

import kafka.common.TopicAndPartition
import kafka.message.MessageAndMetadata
import kafka.serializer.StringDecoder
import org.apache.kafka.clients.consumer.ConsumerConfig
import org.apache.spark.SparkConf
import org.apache.spark.streaming.dstream.InputDStream
import org.apache.spark.streaming.kafka.KafkaCluster.Err
import org.apache.spark.streaming.kafka.{HasOffsetRanges, KafkaCluster, KafkaUtils, OffsetRange}
import org.apache.spark.streaming.{Seconds, StreamingContext}

import scala.collection.immutable.HashMap
import scala.collection.parallel.mutable

/**
 * @description: 低级api,重启后不能恢复消费已经堆积的消息或者换组+指定auto.offset.reset为earliest
 * @time: 2020/6/16 10:44
 * @author: baojinlong
 **/
object LowerKafkaStreaming {
  def main(args: Array[String]): Unit = {
    // 创建SparkConf
    val sparkConf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("wordCount")
    // 创建SteamingContext
    val ssc = new StreamingContext(sparkConf, Seconds(3))
    // 设置kafka参数
    val group: String = "group01"
    val brokers: String = "localhost:9092"
    val topic: String = "test1"
    val deserializationClass: String = "org.apache.kafka.common.serialization.StringDeserializer"

    // 封装Kafka参数
    val kafkaPara: Map[String, String] = Map[String, String](
      // 设置zk配置保存offset
      "zookeeper.connect" -> "localhost:2181",
      ConsumerConfig.GROUP_ID_CONFIG -> group,
      ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG -> brokers,
      ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG -> deserializationClass,
      ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG -> deserializationClass
    )

    // 创建kafka集群对象
    val kafkaCluster = new KafkaCluster(kafkaPara)

    // 获取当前消费offset
    val fromOffset: Map[TopicAndPartition, Long] = getOffset(kafkaCluster, group, topic)

    // 读取kafka数据创建DStream
    val kafkaDstream: InputDStream[String] = KafkaUtils.createDirectStream[String, String, StringDecoder, StringDecoder, String](ssc, kafkaPara, fromOffset, (message: MessageAndMetadata[String, String]) => message.message())
    // 打印数据
    kafkaDstream.print
    // 保存offset
    setOffset(kafkaCluster, group, kafkaDstream)
    // 开启
    ssc.start()
    ssc.awaitTermination()
  }


  /**
   * 保存offset
   *
   * @param kafkaCluster
   * @param kafkaDstream
   */
  def setOffset(kafkaCluster: KafkaCluster, group: String, kafkaDstream: InputDStream[String]): Unit = {
    kafkaDstream.foreachRDD(rdd => {
      // 注意可变的map必须要用var来修饰
      var partitionToLong = new HashMap[TopicAndPartition, Long]
      // 取出rdd的offset
      val offsetRanges: HasOffsetRanges = rdd.asInstanceOf[HasOffsetRanges]
      // 获取所有分区的offsetRange
      val ranges: Array[OffsetRange] = offsetRanges.offsetRanges
      // 遍历数组
      for (range <- ranges) {
        // 遍历数组
        val partition: TopicAndPartition = range.topicAndPartition()
        val tmpOffset: Long = range.untilOffset
        partitionToLong += (partition -> tmpOffset)
      }
      kafkaCluster.setConsumerOffsets(group, partitionToLong)
    })
  }

  def getOffset(kafkaCluster: KafkaCluster, group: String, topic: String): Map[TopicAndPartition, Long] = {
    // 定义最终返回值
    var returnResult = new HashMap[TopicAndPartition, Long]
    // 根据指定的topic获取分区
    val topicAndPartitionsEither: Either[Err, Set[TopicAndPartition]] = kafkaCluster.getPartitions(Set(topic))
    // 判断分区是否存在
    if (topicAndPartitionsEither.isRight) {
      // 分区数据不为空,取出分区信息
      val topicAndPartitions: Set[TopicAndPartition] = topicAndPartitionsEither.right.get
      // 获取消费者消费数据的进度
      val topicAndPartitionToLongEither: Either[Err, Map[TopicAndPartition, Long]] = kafkaCluster.getConsumerOffsets(group, topicAndPartitions)
      // 判断offset是否存在
      if (topicAndPartitionToLongEither.isLeft) {
        // offset不存在(该消费者组未消费过),需要给每个分区都放0
        for (x <- topicAndPartitions) {
          // 注意0L应该给当前分区的最小commitId[使用SimpleConsumer],可以使用低级api来获取,或者更换一个消费者组+指定auto.offset.rest=earliest
          returnResult += (x -> 0L)
        }
      } else {
        // 取出offset
        val value: Map[TopicAndPartition, Long] = topicAndPartitionToLongEither.right.get
        returnResult ++= value
      }
    }
    returnResult
  }

}
